import tensorflow as tf
import numpy as np
import time
from npu_bridge.estimator import npu_ops
from tensorflow.core.protobuf.rewriter_config_pb2 import RewriterConfig

config = tf.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
custom_op.parameter_map["use_off_line"].b = True

custom_op.parameter_map["input_shape"].s = tf.compat.as_bytes("input:1,16,-1,-1")
custom_op.parameter_map["dynamic_dims"].s = tf.compat.as_bytes("200,200;400,400")
custom_op.parameter_map["dynamic_node_type"].i = 1

custom_op.parameter_map["precision_mode"].s = tf.compat.as_bytes("force_fp16")
custom_op.parameter_map["graph_run_mode"].i = 1
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
Graph = tf.Graph()
Graph.as_default()
with tf.gfile.FastGFile("./pixel_shuffle.pb", 'rb') as f:
    Graph_def = tf.GraphDef()
    Graph_def.ParseFromString(f.read())
    tf.import_graph_def(Graph_def, name='')
sess = tf.Session(config=config)
input = sess.graph.get_tensor_by_name('color_input:0')
output = sess.graph.get_tensor_by_name('LeakyRelu:0')
img = np.random.normal(0,1,(1,16,200,200)).astype(np.float16)
result = sess.run([output],feed_dict = {input:img})

